import logging
import pandas as pd
import numpy as np

from coin.strategy.interval.research.dumper import get_dump_root_dir


def load_recent_df():
  logging.info('Loading...')
  dump_root_dir = get_dump_root_dir()
  df = pd.read_pickle(f'{dump_root_dir}/recent/dump.pkl.gz')
  return df


def get_df_from_total(product_name, total_df):
  replace_map = dict()
  for column in total_df[product_name].columns:
    replace_map[column] = column.split('/')[-1]
    if column[0] == 'Y':
      replace_map[column] = 'Y.' + replace_map[column]
  df = total_df[product_name].rename(columns=replace_map)
  df['tdate'] = pd.to_datetime(df['global_timestamp'])
  df = df.set_index('tdate', drop=True)
  df = df.asfreq('1min').reset_index()
  df['global_timestamp'] = df['tdate'].apply(lambda x: x.value)
  df = df.drop('tdate', axis=1)
  # print(df)

  # ffill price columns
  ffill_col = ['b.ask0_p', 'b.bid0_p', 'b.mid_p']
  ffill_col_prefix = ['mid_ret.', 't.vwap.']
  for col in ffill_col:
    df[col] = df[col].ffill()

  for fcol in ffill_col_prefix:
    for col in df.columns:
      if fcol in col:
        df[col] = df[col].ffill()

  return df


def fast_apply(df, alpha_trans_lst, n_cores=16):
  from concurrent.futures import ProcessPoolExecutor
  x_list = list()
  with ProcessPoolExecutor(max_workers=n_cores) as executor:
    futures = list()
    for alpha_trans in alpha_trans_lst:
      futures.append(executor.submit(alpha_trans.eval, df))
    for fut in futures:
      x_list.append(fut.result())

  return x_list


def get_mean_std(x):
  return x.mean(), x.std(ddof=0)


def test1(row):
  return row['A'] + row['B']


def test2(row):
  return row['C'] - row['D']


if __name__ == '__main__':
  length = 1000
  # dates = pd.date_range('20130101', periods=length)
  df = pd.DataFrame(np.random.randn(length, 4), columns=list('ABCD'))
  print(df)
  x = fast_apply(df, [test1, test2])
  print(x)
